from base.vertice import Vertice
from base.graph import Graph
from prim import mst_prim

w = dict()  # dicionario dos pesos das arestas, associa um peso a cada aresta!


def weight(pair):
	if pair in w:
		return w[pair]

vertices = [
	Vertice('A'), Vertice('B'), Vertice('C'), Vertice('D'), Vertice('E'), Vertice('F'), Vertice('G'), Vertice('H'), Vertice('K')
]
g = Graph(vertices)

# A

# adicionando arestas no grafo 1
g.add_edge(vertices[0], vertices[1])
g.add_edge(vertices[1], vertices[0])
# adiciona pesos às arestas
w[(vertices[0], vertices[1])] = 5
w[(vertices[1], vertices[0])] = 5

g.add_edge(vertices[0], vertices[6])
g.add_edge(vertices[6], vertices[0])
w[(vertices[0], vertices[6])] = 10
w[(vertices[6], vertices[0])] = 10

g.add_edge(vertices[0], vertices[3])
g.add_edge(vertices[3], vertices[0])
w[(vertices[0], vertices[3])] = 10
w[(vertices[3], vertices[0])] = 10

g.add_edge(vertices[0], vertices[7])
g.add_edge(vertices[7], vertices[0])
w[(vertices[0], vertices[7])] = 10
w[(vertices[7], vertices[0])] = 10

g.add_edge(vertices[0], vertices[4])
g.add_edge(vertices[4], vertices[0])
w[(vertices[0], vertices[4])] = 5
w[(vertices[4], vertices[0])] = 5

# B

g.add_edge(vertices[1], vertices[5])
g.add_edge(vertices[5], vertices[1])
w[(vertices[1], vertices[5])] = 5
w[(vertices[5], vertices[1])] = 5

g.add_edge(vertices[1], vertices[7])
g.add_edge(vertices[7], vertices[1])
w[(vertices[1], vertices[7])] = 5
w[(vertices[7], vertices[1])] = 5

g.add_edge(vertices[1], vertices[8])
g.add_edge(vertices[8], vertices[1])
w[(vertices[1], vertices[8])] = 10
w[(vertices[8], vertices[1])] = 10

# C

g.add_edge(vertices[2], vertices[3])
g.add_edge(vertices[3], vertices[2])
w[(vertices[2], vertices[3])] = 5
w[(vertices[3], vertices[2])] = 5

g.add_edge(vertices[2], vertices[6])
g.add_edge(vertices[6], vertices[2])
w[(vertices[2], vertices[6])] = 5
w[(vertices[6], vertices[2])] = 5

g.add_edge(vertices[2], vertices[4])
g.add_edge(vertices[4], vertices[2])
w[(vertices[2], vertices[4])] = 10
w[(vertices[4], vertices[2])] = 10

# D

g.add_edge(vertices[3], vertices[6])
g.add_edge(vertices[6], vertices[3])
w[(vertices[3], vertices[6])] = 10
w[(vertices[6], vertices[3])] = 10

g.add_edge(vertices[3], vertices[4])
g.add_edge(vertices[4], vertices[3])
w[(vertices[3], vertices[4])] = 5
w[(vertices[4], vertices[3])] = 5

# E
g.add_edge(vertices[4], vertices[8])
g.add_edge(vertices[8], vertices[4])
w[(vertices[4], vertices[8])] = 10
w[(vertices[8], vertices[4])] = 10

# F

g.add_edge(vertices[5], vertices[6])
g.add_edge(vertices[6], vertices[5])
w[(vertices[5], vertices[6])] = 15
w[(vertices[6], vertices[5])] = 15

# H

g.add_edge(vertices[7], vertices[8])
g.add_edge(vertices[8], vertices[7])
w[(vertices[7], vertices[8])] = 20
w[(vertices[8], vertices[7])] = 20

peso = 0
print('AGM do grafo 1: ')
for u in mst_prim(g, weight, vertices[0]):
	print('(%s, %s)' %(u[0].name, u[1].name))
	peso += u[1].key

print('Peso: ', peso)

print()

vertices = [
	Vertice(1), Vertice(2), Vertice(3), Vertice(4), Vertice(5), Vertice(6), Vertice(7), Vertice(8)
]
g = Graph(vertices)

# 1

# adicionando arestas no grafo
g.add_edge(vertices[0], vertices[1])
g.add_edge(vertices[1], vertices[0])
# adiciona pesos às arestas
w[(vertices[0], vertices[1])] = 8
w[(vertices[1], vertices[0])] = 8

g.add_edge(vertices[0], vertices[6])
g.add_edge(vertices[6], vertices[0])
w[(vertices[0], vertices[6])] = 9
w[(vertices[6], vertices[0])] = 9

g.add_edge(vertices[0], vertices[2])
g.add_edge(vertices[2], vertices[0])
w[(vertices[0], vertices[2])] = 2
w[(vertices[2], vertices[0])] = 2

# 2

g.add_edge(vertices[1], vertices[5])
g.add_edge(vertices[5], vertices[1])
w[(vertices[1], vertices[5])] = 5
w[(vertices[5], vertices[1])] = 5

g.add_edge(vertices[1], vertices[7])
g.add_edge(vertices[7], vertices[1])
w[(vertices[1], vertices[7])] = 6
w[(vertices[7], vertices[1])] = 6

g.add_edge(vertices[1], vertices[4])
g.add_edge(vertices[4], vertices[1])
w[(vertices[1], vertices[4])] = 7
w[(vertices[4], vertices[1])] = 7

g.add_edge(vertices[1], vertices[3])
g.add_edge(vertices[3], vertices[1])
w[(vertices[1], vertices[3])] = 2
w[(vertices[3], vertices[1])] = 2

g.add_edge(vertices[1], vertices[2])
g.add_edge(vertices[2], vertices[1])
w[(vertices[1], vertices[2])] = 3
w[(vertices[2], vertices[1])] = 3

# 3

g.add_edge(vertices[2], vertices[3])
g.add_edge(vertices[3], vertices[2])
w[(vertices[2], vertices[3])] = 9
w[(vertices[3], vertices[2])] = 9

g.add_edge(vertices[2], vertices[6])
g.add_edge(vertices[6], vertices[2])
w[(vertices[2], vertices[6])] = 6
w[(vertices[6], vertices[2])] = 6

# 4

g.add_edge(vertices[3], vertices[4])
g.add_edge(vertices[4], vertices[3])
w[(vertices[3], vertices[4])] = 6
w[(vertices[4], vertices[3])] = 6

g.add_edge(vertices[3], vertices[6])
g.add_edge(vertices[6], vertices[3])
w[(vertices[3], vertices[6])] = 2
w[(vertices[6], vertices[3])] = 2

# 5

g.add_edge(vertices[4], vertices[5])
g.add_edge(vertices[5], vertices[4])
w[(vertices[4], vertices[5])] = 5
w[(vertices[5], vertices[4])] = 5

g.add_edge(vertices[4], vertices[6])
g.add_edge(vertices[6], vertices[4])
w[(vertices[4], vertices[6])] = 4
w[(vertices[6], vertices[4])] = 4

# 6

g.add_edge(vertices[5], vertices[7])
g.add_edge(vertices[7], vertices[5])
w[(vertices[5], vertices[7])] = 4
w[(vertices[7], vertices[5])] = 4

g.add_edge(vertices[5], vertices[6])
g.add_edge(vertices[6], vertices[5])
w[(vertices[5], vertices[6])] = 6
w[(vertices[6], vertices[5])] = 6

# 7

g.add_edge(vertices[6], vertices[7])
g.add_edge(vertices[7], vertices[6])
w[(vertices[6], vertices[7])] = 3
w[(vertices[7], vertices[6])] = 3

peso = 0
print('AGM do grafo 2: ')
for u in mst_prim(g, weight, vertices[0]):
	print('(%s, %s)' %(u[0].name, u[1].name))
	peso += u[1].key

print('Peso: ', peso)
